/*
 * <author>Hankcs</author>
 * <email>me@hankcs.com</email>
 * <create-date>2017-11-02 12:09</create-date>
 *
 * <copyright file="Demo.java" company="码农场">
 * Copyright (c) 2017, 码农场. All Right Reserved, http://www.hankcs.com/
 * This source is subject to Hankcs. Please contact Hankcs to get more information.
 * </copyright>
 */
package com.hankcs.demo;

import com.hankcs.hanlp.corpus.MSR;
import com.hankcs.hanlp.corpus.io.IOUtil;
import com.hankcs.hanlp.mining.word2vec.DocVectorModel;
import com.hankcs.hanlp.mining.word2vec.Word2VecTrainer;
import com.hankcs.hanlp.mining.word2vec.WordVectorModel;
import com.hankcs.hanlp.utility.TestUtility;
import org.shj.util.PropertyUtil;

import java.io.IOException;
import java.util.Map;

/**
 * 演示词向量的训练与应用
 *
 * @author hankcs
 */
public class DemoWord2Vec
{
    private static final String TRAIN_FILE_NAME = PropertyUtil.getRootDataFolder() + "/" + MSR.TRAIN_PATH;
    private static final String MODEL_FILE_NAME = PropertyUtil.getRootDataFolder() + "/" + "data/test/word2vec.txt";

    public static void main(String[] args) throws IOException
    {
        WordVectorModel wordVectorModel = trainOrLoadModel();
        // 文档向量
        DocVectorModel docVectorModel = new DocVectorModel(wordVectorModel);

        /*printNearest("上海", wordVectorModel);
        printNearest("美丽", wordVectorModel);
        printNearest("购买", wordVectorModel);
        System.out.println(wordVectorModel.similarity("上海", "广州"));
        System.out.println(wordVectorModel.analogy("日本", "自民党", "共和党"));


        String[] documents = new String[]{
            "山东苹果丰收",
            "农民在江苏种水稻",
            "奥运会女排夺冠",
            "世界锦标赛胜出",
            "中国足球失败",
        };

        System.out.println(String.format("【%s】与【%s】的相似度是：%f", documents[0], documents[1], docVectorModel.similarity(documents[0], documents[1])));
        System.out.println(String.format("【%s】与【%s】的相似度是：%f", documents[0], documents[4], docVectorModel.similarity(documents[0], documents[4])));


        for (int i = 0; i < documents.length; i++)
        {
            docVectorModel.addDocument(i, documents[i]);
        }

        printNearestDocument("体育", documents, docVectorModel);
        printNearestDocument("农业", documents, docVectorModel);
        printNearestDocument("我要看比赛", documents, docVectorModel);
        printNearestDocument("要不做饭吧", documents, docVectorModel);*/


        System.out.println("================== ");
        System.out.println(docVectorModel.similarity("朱月珍家使用的洗衣机接通老落水管导致堵塞不畅（小区改造时已经排了新落水管），造成徐泉生家里漏水。",
                "郑红侠与朱礼需因房子出租后，室内物品损坏赔偿而发生争执。"));
        System.out.println(docVectorModel.similarity("朱月珍家使用的洗衣机接通老落水管导致堵塞不畅（小区改造时已经排了新落水管），造成徐泉生家里漏水。",
                "5月4日，金色家园48-301蒋利华家厨房进水管漏水，造成101室周莉家厨房间及客厅部分积水，造成损失。双方就赔偿事宜产生争议。"));
        System.out.println(docVectorModel.similarity("朱月珍家使用的洗衣机接通老落水管导致堵塞不畅（小区改造时已经排了新落水管），造成徐泉生家里漏水。",
                "中国足球再一次未能进入世界杯。让我们等待他们下一次的表现。"));
        System.out.println(docVectorModel.similarity("朱月珍家使用的洗衣机接通老落水管导致堵塞不畅（小区改造时已经排了新落水管），造成徐泉生家里漏水。",
                "小明今天晚上吃了荷包蛋、排骨和一大碗米饭"));
        System.out.println(docVectorModel.similarity("朱月珍家使用的洗衣机接通老落水管导致堵塞不畅（小区改造时已经排了新落水管），造成徐泉生家里漏水。",
                "虞毛毛与方阿多是再婚夫妻，虞毛毛在婚前自行购买了一套商品房（枫桥镇卫生院北商品房1-204,1997年12月购买，面积为57平方米）。虞毛毛现对该套房屋买卖后的分配以及百年后遗产分割产生异议。"));

    }

    static void printNearest(String word, WordVectorModel model)
    {
        System.out.printf("\n            Word     Cosine\n--------------------------------------\n");
        for (Map.Entry<String, Float> entry : model.nearest(word))
        {
            System.out.printf("%12s\t\t%f\n", entry.getKey(), entry.getValue());
        }
    }

    static void printNearestDocument(String document, String[] documents, DocVectorModel model)
    {
        printHeader(document);
        for (Map.Entry<Integer, Float> entry : model.nearest(document))
        {
            System.out.printf("%12s\t\t%f\n", documents[entry.getKey()], entry.getValue());
        }
    }

    private static void printHeader(String query)
    {
        System.out.printf("\n%12s          Cosine\n----------------------------------------------\n", query);
    }

    static WordVectorModel trainOrLoadModel() throws IOException
    {
        if (!IOUtil.isFileExisted( MODEL_FILE_NAME))
        {
            if (!IOUtil.isFileExisted(TRAIN_FILE_NAME))
            {
                System.err.println("语料不存在，请阅读文档了解语料获取与格式：https://github.com/hankcs/HanLP/wiki/word2vec");
                System.exit(1);
            }
            Word2VecTrainer trainerBuilder = new Word2VecTrainer();
            return trainerBuilder.train(TRAIN_FILE_NAME, MODEL_FILE_NAME);
        }

        return loadModel();
    }

    static WordVectorModel loadModel() throws IOException
    {
        return new WordVectorModel(MODEL_FILE_NAME);
    }
}
